from datetime import datetime
from jqdata import *
from jqlib.technical_analysis import *
import pandas as pd


# 统计CCI技术指标结果数据
def statistic_CCI(stock, start_date, end_date):
    # 当cci上穿+100时，买入，当cci下穿+100时，卖出

    '''
    构造dataframe:信号日期、信号类型（buy/sell）、当天收盘价、信号出现后第1/3/5/20天收盘价

    date  type  close_current  close_1  close_3  close_5  close_20
    '''
    statistic_df = pd.DataFrame(columns=["date", "type", "close_current", "close_1", "close_3", "close_5", "close_20"])

    # 获取所有交易日，以便获取信号出现后未来20个交易日的价格
    trade_days = get_trade_days(start_date=start_date, end_date=end_date)

    # 昨日CCI
    pre_cci = NaN
    # 今日CCI
    cci = NaN
    for i in range(0, len(trade_days)):
        # 如果往后不足20个交易日，比如昨天刚出现的信号，后面还没数据呢，暂时先不统计进去
        price_end_date_index = i + 20
        if price_end_date_index >= len(trade_days):
            break

        check_date = trade_days[i]
        cci = CCI(stock, check_date=check_date, N=14)[stock]

        # 信号类型
        signal_type = None
        # 买入
        if pre_cci < 100 and cci > 100:
            signal_type = "buy"
        elif pre_cci > 100 and cci < 100:
            signal_type = "sell"

        if not signal_type is None:
            price_end_date = trade_days[price_end_date_index]
            price_close = get_price(stock, end_date=price_end_date, count=21, frequency='1d',
                                    fields='close')['close']
            row = [
                check_date,
                signal_type,
                price_close[0],
                price_close[1],
                price_close[3],
                price_close[5],
                price_close[20]
            ]
            statistic_df.loc[len(statistic_df.index)] = row

        pre_cci = cci
    return statistic_df


# 根据上面统计出的信号出现后的价格数据，进一步统计成功率
def statistic_win_rate(statistic_price_df):
    '''
    构造dataframe:
    信号类型（buy/sell）、信号出现1/3/5/20天后成功率、信号出现1/3/5/20天后平均涨跌幅(%)
    '''
    # 每次信号出现后1/3/5/20天股价涨跌幅（%， 2位小数）
    statistic_price_df['chg_pct_1'] = round(
        (statistic_price_df['close_1'] - statistic_price_df['close_current']) / statistic_price_df[
            'close_current'] * 100, 2)
    statistic_price_df['chg_pct_3'] = round(
        (statistic_price_df['close_3'] - statistic_price_df['close_current']) / statistic_price_df[
            'close_current'] * 100, 2)
    statistic_price_df['chg_pct_5'] = round(
        (statistic_price_df['close_5'] - statistic_price_df['close_current']) / statistic_price_df[
            'close_current'] * 100, 2)
    statistic_price_df['chg_pct_20'] = round(
        (statistic_price_df['close_20'] - statistic_price_df['close_current']) / statistic_price_df[
            'close_current'] * 100, 2)

    win_rate_df = pd.DataFrame(columns=['信号类型', '出现次数',
                                        '1日胜率（%）', '1日平均涨跌幅（%）',
                                        '3日胜率（%）', '3日平均涨跌幅（%）',
                                        '5日胜率（%）', '5日平均涨跌幅（%）',
                                        '20日胜率（%）', '20日平均涨跌幅（%）'])
    # buy info
    buy_df = statistic_price_df[statistic_price_df['type'] == 'buy']
    buy_count = buy_df.shape[0]
    # 涨跌幅 > 0 算win
    buy_win_rate_1 = round(buy_df[buy_df['chg_pct_1'] > 0].shape[0] / buy_count * 100, 2)
    avg_chg_pct_1 = round(buy_df['chg_pct_1'].mean(), 2)
    buy_win_rate_3 = round(buy_df[buy_df['chg_pct_3'] > 0].shape[0] / buy_count * 100, 2)
    avg_chg_pct_3 = round(buy_df['chg_pct_3'].mean(), 2)
    buy_win_rate_5 = round(buy_df[buy_df['chg_pct_5'] > 0].shape[0] / buy_count * 100, 2)
    avg_chg_pct_5 = round(buy_df['chg_pct_5'].mean(), 2)
    buy_win_rate_20 = round(buy_df[buy_df['chg_pct_20'] > 0].shape[0] / buy_count * 100, 2)
    avg_chg_pct_20 = round(buy_df['chg_pct_20'].mean(), 2)

    buy_row = ['买', buy_count,
               buy_win_rate_1, avg_chg_pct_1,
               buy_win_rate_3, avg_chg_pct_3,
               buy_win_rate_5, avg_chg_pct_5,
               buy_win_rate_20, avg_chg_pct_20]

    win_rate_df.loc[0] = buy_row

    # sell info
    sell_df = statistic_price_df[statistic_price_df['type'] == 'sell']
    sell_count = sell_df.shape[0]
    # 涨跌幅 < 0 算win
    sell_win_rate_1 = round(sell_df[sell_df['chg_pct_1'] < 0].shape[0] / sell_count * 100, 2)
    avg_chg_pct_1 = round(sell_df['chg_pct_1'].mean(), 2)
    sell_win_rate_3 = round(sell_df[sell_df['chg_pct_3'] < 0].shape[0] / sell_count * 100, 2)
    avg_chg_pct_3 = round(sell_df['chg_pct_3'].mean(), 2)
    sell_win_rate_5 = round(sell_df[sell_df['chg_pct_5'] < 0].shape[0] / sell_count * 100, 2)
    avg_chg_pct_5 = round(sell_df['chg_pct_5'].mean(), 2)
    sell_win_rate_20 = round(sell_df[sell_df['chg_pct_20'] < 0].shape[0] / sell_count * 100, 2)
    avg_chg_pct_20 = round(sell_df['chg_pct_20'].mean(), 2)

    sell_row = ['卖', sell_count,
                sell_win_rate_1, avg_chg_pct_1,
                sell_win_rate_3, avg_chg_pct_3,
                sell_win_rate_5, avg_chg_pct_5,
                sell_win_rate_20, avg_chg_pct_20]

    win_rate_df.loc[1] = sell_row

    return win_rate_df

price_df = statistic_CCI("600036.XSHG", "2013-01-01", "2022-09-30")
win_rate_df = statistic_win_rate(price_df)
# print(win_rate_df)